package com.shujia.mllib

import org.apache.spark.ml.linalg
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.{DataFrame, DataFrameReader, SaveMode, SparkSession}

object Demo4MakeDataImage {
  def main(args: Array[String]): Unit = {

    val spark: SparkSession = SparkSession
      .builder()
      .master("local[8]")
      .appName("image")
      .getOrCreate()

    import spark.implicits._

    val data: DataFrame = spark
      .read
      .format("image")
      .load("D:\\课件\\机器学习数据\\手写数字\\train")

    data.printSchema()

    val nameDF: DataFrame = data
      .select($"image.origin" as "origin", $"image.data" as "data")
      .as[(String, Array[Byte])]
      .rdd
      .map {
        case (origin: String, dataArr: Array[Byte]) =>

          //获取图片名
          val name: String = origin.split("/").last


          val arr: Array[Double] = dataArr
            .map(_.toInt)
            //将数据进行归一化，变成0或者1
            .map(i => {
            if (i < 0) {
              1.0
            } else {
              0.0
            }
          })

          //将数据转换成向量
          val vector: linalg.Vector = Vectors.dense(arr)


          (name, vector)
      }
      .toDF("name", "features")


    //加载图片对应数字的文件
    val labelDF: DataFrame = spark
      .read
      .format("csv")
      .option("sep", " ")
      .schema("name STRING,label DOUBLE")
      .load("D:\\课件\\机器学习数据\\手写数字\\train.txt")


    //关联标签，将文件名缓存数字
    nameDF
      .join(labelDF, "name")
      .select($"label", $"features")
      .write
      .format("libsvm")
      .mode(SaveMode.Overwrite)
      .save("Spark/data/image_data")


  }

}
